# !/bin/env python
# -*- encoding: utf-8 -*-

import os
# https://www.nowcoder.com/discuss/554260171740958720?fromPut=jj-github&urlSource=extension-api
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility


class MilvusBase(object):
    def __init__(self, **kwargs):
        self.config = kwargs
        self.collection = None
        connections.connect(host=self.config["milvus_host"], port=self.config["milvus_port"])

    def create_question_collection(self, collection_name):
        if not self.has_collection(collection_name):
            field1 = FieldSchema(name="id", dtype=DataType.INT64, descrition="int64", is_primary=True, auto_id=True)
            field2 = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, descrition="float vector",
                                 dim=self.config["VECTOR_DIM"], is_primary=False)
            field3 = FieldSchema(name="question", dtype=DataType.VARCHAR, description="question", is_primary=False,
                                 max_length=65530)
            field4 = FieldSchema(name="answer", dtype=DataType.VARCHAR, description="answer", is_primary=False,
                                 max_length=65530)
            schema = CollectionSchema(fields=[field1, field2, field3, field4], description="collection description")
            self.collection = Collection(name=collection_name, schema=schema)
            self.create_index(collection_name)




    def create_text_collection(self, collection_name):
        if not self.has_collection(collection_name):
            field1 = FieldSchema(name="id", dtype=DataType.INT64, descrition="int64", is_primary=True, auto_id=True)
            field2 = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, descrition="float vector",
                                 dim=self.config["VECTOR_DIM"], is_primary=False)
            field3 = FieldSchema(name="text", dtype=DataType.VARCHAR, description="text", is_primary=False,
                                 max_length=65530)
            schema = CollectionSchema(fields=[field1, field2, field3], description="collection description")
            self.collection = Collection(name=collection_name, schema=schema)
            self.create_index(collection_name)

    def set_collection(self, collection_name):
        if self.has_collection(collection_name):
            self.collection = Collection(name=collection_name)

    def has_collection(self, collection_name):
        return utility.has_collection(collection_name)

    def insert_question(self, collection_name, vectors, table_name, table_fields):
        # insert qs to milvus collection
        self.set_collection(collection_name)
        data = [[vectors], [table_name], [table_fields]]
        mr = self.collection.insert(data)
        ids = mr.primary_keys
        self.collection.load()
        return ids

    def insert_text(self, collection_name, vectors, text):
        # insert text to milvus collection
        self.set_collection(collection_name)
        data = [[vectors], [text]]
        mr = self.collection.insert(data)
        ids = mr.primary_keys
        self.collection.load()
        return ids

    def create_index(self, collection_name):
        # Create IVF_FLAT index on milvus collection
        default_index = {"index_type": "HNSW", "metric_type": self.config["METRIC_TYPE"],
                         "params": {"M": 32, "efConstruction": 128}}
        # 上面的M原始值为8，现在改为32，M：建表期间每个向量的边数目，M越大，内存消耗越高，高维度的数据集下查询的性能会更好，建议范围8-32
        # efConstruction:原始值为64，现在改为128 控制索引时间和索引准确度，越大，构建的索引越长，查询精度越高
        status = self.collection.create_index(field_name="embedding", index_params=default_index)

    def search_question(self, collection_name, vectors, top_k):
        # Search vector in milvus collection
        self.set_collection(collection_name)
        self.collection.load()
        search_params = {"metric_type": self.config["METRIC_TYPE"], "params": {"M": 8, "efConstruction": 64}}
        res = self.collection.search(vectors, anns_field="embedding", param=search_params, limit=top_k,
                                     output_fields=['question', 'answer'])
        # import pdb;pdb.set_trace()
        # print(res[0])
        # print(res)
        # return res[0][0].entity.table_name, res[0][0].entity.table_fields
        table_name_list = []
        table_fields_list = []
        dis_list = []
        for i in range(len(res[0])):
            table_name_list.append(res[0][i].entity.question)
            table_fields_list.append(res[0][i].entity.answer)
            dis_list.append(res[0][i].distance)

        return table_name_list, table_fields_list, dis_list

    def search_text(self, collection_name, vectors, top_k):
        # Search vector in milvus collection
        self.set_collection(collection_name)
        self.collection.load()
        search_params = {"metric_type": self.config["METRIC_TYPE"], "params": {"M": 8, "efConstruction": 64}}
        res = self.collection.search(vectors, anns_field="embedding", param=search_params, limit=top_k,
                                     output_fields=['text'])

        text_list = []
        dis_list = []
        for i in range(len(res[0])):
            text_list.append(res[0][i].entity.text)
            dis_list.append(res[0][i].distance)

        return text_list, dis_list

    def delete_vectors(self, collection_name):
        self.set_collection(collection_name)
        expr = "id in [446555624057930045,446555624057931601]"
        self.collection.delete(expr)

